///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file indexing.cpp
///
/// \brief   implementation for functions associated with configuration indexing
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include <vector>
#include <numeric>
#include <algorithm>
#include <cassert>

#include <src/indexing.h>

#include <iostream>
#include <cstdio>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   computes the binomial coefficient
///
/// \param[in]     n        number of things
/// \param[in]     m        how many things to take at a time
///
/// \return the number of ways n things can be taken m at a time
///
///////////////////////////////////////////////////////////////////////////////////////////////////
int eds::binom_coeff(int n, int m) {
  if (n < 0 || m < 0 || m > n) return 0.0;
  double retval = 1.0;
  while (m > 0) retval = ( retval * (n--) ) / (m--);
  return int(retval+0.5);
}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Returns the number of solutions to the equation
///          x(1) + x(2) + ... + x(n) = r
///          when the variables x(i) are constrained to be
///          integers in the range (0, 1, 2, ..., k)
///
/// \param[in]     n        number of variables
/// \param[in]     r        sum of varaibles
/// \param[in]     k        range of each variable
/// \param[out]    work     integer workspace, size >= k+1
///
///////////////////////////////////////////////////////////////////////////////////////////////////
int eds::n_integer_solutions(const int n, const int r, const int k, int * work) {

  //// prepare a vector to hold the number of variables having each allowed value
  //static std::vector<int> n_with_value;
  //if (n_with_value.size() < k+1)
  //  n_with_value.resize(k+1);

  // initialize an array to hold the number of variables having each allowed value
  int * const n_with_value = work;
  work += (k+1);
  for (int i = 0; i <= k; i++)
    n_with_value[i] = 0;

  // initialize the return value
  int retval = 0;

  // Loop over all possible distributions of variables among the values.
  // Note that we do not directly track of how many variables are equal to zero,
  // as this is known by how many variables take on other values.
  while (true) {

    // compute the number of nonzero variables
    int n_nonzero = 0;
    for (int i = 1; i <= k; i++)
      n_nonzero += n_with_value[i];

    // compute the sum of the variables
    int sum = 0;
    for (int i = 1; i <= k; i++)
      sum += i * n_with_value[i];

    // if this distribution solves the equation, count how many ways it can occur
    if (sum == r && n_nonzero <= n) {

      // determine how many variables are nonzero
      int t = 0;
      for (int i = 1; i <= k; i++)
        t += n_with_value[i];

      // count how many ways the variables can satisfy this distribution
      int occurances = eds::binom_coeff(n, t);
      for (int i = 1; i < k; i++) {
        occurances *= eds::binom_coeff(t, n_with_value[i]);
        t -= n_with_value[i]; // t is now equal to the number of variables greater than i
      }

      // record how many ways the variables satisfy this distribution
      retval += occurances;

    }

    // increment to the next distribution of variables
    int p;
    for (p = k; p > 0; p--)
      if (++n_with_value[p] > n)
        n_with_value[p] = 0;
      else
        break;

    // stop iterating if all distributions have been processed
    if (p == 0) break;

  }

  // return the result
  return retval;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   computes the compound index for the supplied occupation vector
///
/// \param[in]     n         the number of sites
/// \param[in]     cps       the number of configurations per site
/// \param[in]     t         the total occupation, i.e. the sum of all sites' occupations
/// \param[in]     offsets   array containing offset information for different occupation patterns
/// \param[in]     spacing   array containing spacing information for different occupation patterns
/// \param[in]     z_lookup  array used for looking up the appropriate z array
/// \param[in]     z_arrays  contains all the z arrays
/// \param[in]     occ_vec   the occupation number vector
/// \param[out]    workspace integer workspace, size >= 5*cps + 2
///
///////////////////////////////////////////////////////////////////////////////////////////////////
int eds::index_from_occ_vec(const int n,
                            const int cps,
                            const int t,
                            const int * const offsets,
                            const int * const spacing,
                            const int * const z_lookup,
                            const int * const z_arrays,
                            const int * const occ_vec,
                            int * workspace) {

  // get space to store the number of occurences for each occupation
  int * const occurances = workspace;
  workspace = workspace + cps;
  for (int i = 0; i < cps; i++)
    occurances[i] = 0;

  // get the number of occurances for each occupation
  {
    int i = 0;
    for ( ; i + 3 < n; i+=4) {
      occurances[occ_vec[i+0]]++;
      occurances[occ_vec[i+1]]++;
      occurances[occ_vec[i+2]]++;
      occurances[occ_vec[i+3]]++;
    }
    for ( ; i < n; i++)
      occurances[occ_vec[i+0]]++;
    assert( std::accumulate(occurances, occurances+cps, 0) == n );
  }

  // compute the size of the testing arrays
  const int test_array_size = 2*cps+1;

  // get an array used for testing that a value is zero
  int * const eq_zero_array = workspace;
  workspace = workspace + test_array_size;
  for (int i = 0; i < test_array_size; i++)
    eq_zero_array[i] = 0;
  eq_zero_array[cps+1] = 1;
  const int * const eq_zero = eq_zero_array + (cps+1);

  // get an array used for testing that a value is <= zero
  int * const le_zero_array = workspace;
  workspace = workspace + test_array_size;
  for (int i = 0; i <= cps+1; i++)
    le_zero_array[i] = 1;
  for (int i = cps+2; i < test_array_size; i++)
    le_zero_array[i] = 0;
  const int * const le_zero = le_zero_array + (cps+1);

  // initialize the compound index
  int cmpd = 0;

  // initialize the number of sites remaining
  int rem_sites = n;

  // initialize the remaining occupation
  int rem_occ = t;

  // compute the compound index
  for (int occ = cps-1; occ > 0; occ--) {

    // get the number of occurances of this occupation
    const int p = occurances[occ];

    // get the appropriate z pointer
    const int * const z_ptr = z_arrays + z_lookup[ (n+1)*rem_sites + p ];

    // add the offset
    if (occ > 1) {
      const int off = offsets[ occ*(t+1)*(n+1)*(n+1) + rem_sites*(t+1)*(n+1) + rem_occ*(n+1) + p ];
      //std::cout << " offset[" << occ << "] = " << off << std::endl;
      cmpd += off;
    }

    // compute the z index
    int z_ind = 0;
    {

      int i = 0;
      int count = 0;
      int pos = 0;
      int x, diff;

      for ( ; i + 2 < n; i+=3) {

        diff = occ_vec[i+0] - occ;
        x = eq_zero[diff];
        z_ind += x * z_ptr[ rem_sites*count + pos ];
        count += x;
        pos += le_zero[diff];

        diff = occ_vec[i+1] - occ;
        x = eq_zero[diff];
        z_ind += x * z_ptr[ rem_sites*count + pos ];
        count += x;
        pos += le_zero[diff];

        diff = occ_vec[i+2] - occ;
        x = eq_zero[diff];
        z_ind += x * z_ptr[ rem_sites*count + pos ];
        count += x;
        pos += le_zero[diff];

      }

      for ( ; i < n; i++) {

        diff = occ_vec[i+0] - occ;
        x = eq_zero[diff];
        z_ind += x * z_ptr[ rem_sites*count + pos ];
        count += x;
        pos += le_zero[diff];

      }

      assert( count == p );
      //assert( pos == rem_sites );

    }

    // add the z index times the spacing
    const int sp = ( occ == 1 ? 1 : spacing[ (occ-1)*(t+1)*(n+1) + (rem_sites-p)*(t+1) + rem_occ-p*occ ] );
    cmpd += z_ind * sp;
    //std::cout << "  z_ind[" << occ << "] = " << z_ind << std::endl;
    //std::cout << "spacing[" << occ << "] = " << sp << std::endl;

    // update the number of rem_sites sites
    rem_sites -= p;

    // update the remaining occupation
    rem_occ -= p*occ;

  }

  //return the compound index
  return cmpd;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Prepares the spacing and offset arrays used in compound indexing
///
/// \param[in]     n         the number of sites
/// \param[in]     cps       the number of configurations per site
/// \param[in]     t         the total occupation, i.e. the sum of all sites' occupations
/// \param[in,out] offsets   array to contain offset information for different occupation patterns
///                          on entry, size must be (t+1)*(n+1)*(n+1)*cps
/// \param[in,out] spacing   array to contain spacing information for different occupation patterns
///                          on entry, size must be (t+1)*(n+1)*cps
/// \param[out]    work     integer workspace, size >= cps
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void eds::prepare_spacing_and_offsets(const int n,
                                      const int cps,
                                      const int t,
                                      int * const offsets,
                                      int * const spacing,
                                      int * work) {

  // get spacings, which are the number of arrangements for a given particle rank, # sites, and total occupation
  // here  0 < k < cps
  for (int k = 1; k < cps; k++) // number of particle types
  for (int m = 0; m <= n; m++)  // number of sites
  for (int s = 0; s <= t; s++)  // total occupation
    spacing[ k*(t+1)*(n+1) + m*(t+1) + s ] = eds::n_integer_solutions(m, s, k, work);

  // get offsets
  //std::printf("%12i\n", eds::n_integer_solutions(2, 2, 1, work));
  //std::printf("%9s   %9s   %9s   %9s   %12s   %12s", "part type", "# sites", "total occ", "# part", "offset", "spacing");
  //std::cout << std::endl;
  for (int k = 2; k < cps; k++)             // particle type
  for (int m = 0; m <= n; m++)              // number of sites
  for (int s = 0; s <= t; s++)              // total occupation
  for (int q = 0; q <= m && q*k <= s; q++)  // number of particles of this type
  {
    const int x = k*(t+1)*(n+1)*(n+1) + m*(t+1)*(n+1) + s*(n+1) + q;
    int sp = 0;
    if (q == 0) {
      offsets[x] = 0;
    } else {
      const int bc = eds::binom_coeff(m, q-1);
      sp = spacing[ (k-1)*(t+1)*(n+1) + (m-q+1)*(t+1) + s-(q-1)*k ];
      offsets[x] = offsets[x-1] + bc * sp;
    }
    //if ( s - q*k <= (k-1) * (m-q) && s <= m * k ) {
    //  std::printf("%9i   %9i   %9i   %9i   %12i   %12i", k, m, s, q, offsets[x], sp);
    //  std::cout << std::endl;
    //}
  }

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Returns the size of the z array storage for a given lattice length.
///
/// \param[in]     n         the number of sites in the lattice
///
///////////////////////////////////////////////////////////////////////////////////////////////////
int eds::z_array_size(const int n) {

  // initialize the return value
  int retval = 0;

  // loop over the possible combinations of sites and particles
  for (int s = 0; s <= n; s++)  // sites
  for (int p = 0; p <= s; p++) // particles
  {

    // add the number of elements needed to the total
    retval += std::max(s*p, 1);

  }

  // return the total
  return retval;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Prepares the z array used in computing compound indeces for configurations of M
///          identical particles in N identical boxes.
///
/// \param[in]     n         the number of sites
/// \param[in,out] z_array   array to contain the z arrays
///                          on entry, size must be eds::z_array_size(n)
/// \param[in,out] z_lookup  array used to find the position of a given z array in z_array
///                          on entry, size must be (n+1)*(n+1)
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void eds::prepare_z_array(const int n, int * const z_array, int * const z_lookup) {

  // get a variable to track the positions of the arrays
  int pos = 0;

  // loop over the possible combinations of sites and particles
  for (int s = 0; s <= n; s++) // sites
  for (int p = 0; p <= s; p++) // particles
  {

    // get a pointer to this z array
    int * const z_ptr = z_array + pos;

    // set the lookup position
    z_lookup[s*(n+1)+p] = pos;

    // increment the position for the next array
    pos += std::max(s*p, 1);

    // if there are no particles, the setup is trivial
    if (p == 0) {
      z_ptr[0] = 0;
      continue;
    }

    // compute the array element for each (particle number, position) pair
    for (int k = 0; k < p; k++)    // particle number
    for (int l = 0; l < s; l++) {  // position
      const int M = s;
      const int N = p;
      const int K = k+1;
      const int L = l+1;
      int z = 0;
      for (int m = M - L + 1; m <= M - K; m++)
        z += eds::binom_coeff(m, N-K) - eds::binom_coeff(m-1, N-K-1);
      z_ptr[k*s+l] = z;
      assert( z_lookup[s*n+p] + k*s+l < pos );
    }

  }

}
